﻿//------------------------------------------------------------------------------
// <copyright file="RpcSocketTransport.cs" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------

namespace Microsoft.Internal.GamesTest.Rpc.Client.Transport
{
    using System;
    using System.Collections.Concurrent;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Globalization;
    using System.IO;
    using System.Net;
    using System.Net.Sockets;
    using System.Reflection;
    using System.Runtime.Serialization;
    using System.Runtime.Serialization.Json;
    using System.Text;
    using System.Threading;
    using System.Threading.Tasks;

    /// <summary>
    /// Transport to connect using sockets. Used for Windows targets.
    /// </summary>
    internal class RpcSocketTransport : DisposableObject
    {
        private const int DefaultRpcServerPort = 4600;
        private const int MaximumOperationTimeoutMilliseconds = unchecked((int)0x7FFFFFFF);

        private const string IncomingSerializedMessagesFailure = "Failed to parse incoming RpcSocketTransportMessage. See inner exception for details.";

        private readonly DataContractJsonSerializer socketTransportMessageSerializer;

        private Socket socket;
        private ReaderWriterLockSlim socketLock = new ReaderWriterLockSlim();
        private object networkStreamLock = new object();

        private IPAddress ipAddress;
        private int portNumber;
        private Guid serverId;
        private TimeSpan operationTimeout;
        private ReaderWriterLockSlim operationTimeOutLock = new ReaderWriterLockSlim();

        /// <summary>
        /// Initializes a new instance of the RpcSocketTransport class.
        /// </summary>
        /// <param name="ipAddress">The IP Address of the target.</param>
        /// <param name="serverId">The Id of the RPC server to use on the target.</param>
        public RpcSocketTransport(IPAddress ipAddress, Guid serverId)
        {
            this.socket = null;
            this.ipAddress = ipAddress;
            this.portNumber = DefaultRpcServerPort;
            this.serverId = serverId;
            this.OperationTimeout = new TimeSpan();

            this.socketTransportMessageSerializer = new DataContractJsonSerializer(typeof(RpcSocketTransportMessage));
        }

        /// <summary>
        /// Initializes a new instance of the RpcSocketTransport class with port number.
        /// </summary>
        /// <param name="ipAddress">The IP Address of the target.</param>
        /// <param name="port">The Port number of the target.</param>
        /// <param name="serverId">The Id of the RPC server to use on the target.</param>
        public RpcSocketTransport(IPAddress ipAddress, int port, Guid serverId)
        {
            this.socket = null;
            this.ipAddress = ipAddress;
            this.portNumber = port;
            this.serverId = serverId;
            this.OperationTimeout = new TimeSpan();

            this.socketTransportMessageSerializer = new DataContractJsonSerializer(typeof(RpcSocketTransportMessage));
        }

        /// <summary>
        /// Gets the default port for connections to Rpc server.
        /// </summary>
        public static int DefaultPort
        {
            get
            {
                return DefaultRpcServerPort;
            }
        }

        /// <summary>
        /// Gets the port for this connection.
        /// </summary>
        public int Port
        {
            get
            {
                return this.portNumber;
            }
        }

        /// <summary>
        /// Gets the server id for this connection.
        /// </summary>
        public Guid ServerId
        {
            get
            {
                return this.serverId;
            }
        }

        /// <summary>
        /// Gets or sets the timeout for all operations. OperationTimeout can only be set to System.Threading.Timeout.InfiniteTimeSpan or a TimeSpan value with total milliseconds greater or equal to 0 and less than 2147483648.
        /// </summary>
        public TimeSpan OperationTimeout
        {
            get
            {
                try
                {
                    this.operationTimeOutLock.EnterReadLock();

                    return this.operationTimeout;
                }
                finally
                {
                    this.operationTimeOutLock.ExitReadLock();
                }
            }

            set
            {
                try
                {
                    this.operationTimeOutLock.EnterWriteLock();

                    ThrowIfInvalidTimeout(value, "value");

                    this.operationTimeout = value;
                }
                finally
                {
                    this.operationTimeOutLock.ExitWriteLock();
                }
            }
        }

        /// <summary>
        /// Gets a value indicating whether the transport is connected to the target server.
        /// </summary>
        public bool IsConnected
        {
            get
            {
                bool connected = false;

                try
                {
                    this.socketLock.EnterReadLock();

                    if (this.socket == null)
                    {
                        return false;
                    }

                    connected = this.socket.Connected;
                }
                finally
                {
                    this.socketLock.ExitReadLock();
                }

                return connected;
            }
        }

        /// <summary>
        /// Gets the internal network stream. Marked as protected to allow unit tests to shim it out.
        /// </summary>
        protected NetworkStream NetworkStream { get; private set; }

        /// <summary>
        /// Helper method to throw on invalid timeout values.
        /// </summary>
        /// <param name="value">The timeout value to validate.</param>
        /// <param name="valueName">The value name to use in the error message.</param>
        public static void ThrowIfInvalidTimeout(TimeSpan? value, string valueName)
        {
            if (value != null)
            {
                ThrowIfInvalidTimeout(value.Value, valueName);
            }
        }

        /// <summary>
        /// Helper method to throw on invalid timeout values.
        /// </summary>
        /// <param name="value">The timeout value to validate.</param>
        /// <param name="valueName">The value name to use in the error message.</param>
        public static void ThrowIfInvalidTimeout(TimeSpan value, string valueName)
        {
            if (value == null)
            {
                return;
            }

            double totalMilliseconds = value.TotalMilliseconds;

            if (value != Timeout.InfiniteTimeSpan && (totalMilliseconds < 0 || totalMilliseconds > MaximumOperationTimeoutMilliseconds))
            {
                throw new ArgumentOutOfRangeException(
                    valueName,
                    value,
                    string.Format(CultureInfo.CurrentCulture, "{0} can only be set to System.Threading.Timeout.InfiniteTimeSpan or a TimeSpan value with total milliseconds greater or equal to 0 and less than 0x{1:x}.", valueName, MaximumOperationTimeoutMilliseconds));
            }
        }

        /// <summary>
        /// Send a request message to the RPC Server.  The response will be included in the IEnumerable
        /// returned by the GetIncomingSerializedMessages method.
        /// </summary>
        /// <param name="serializedMessage">A serialized RpcMessage.</param>
        public void SendRequest(string serializedMessage)
        {
            this.SendRequest(serializedMessage, this.OperationTimeout);
        }

        /// <summary>
        /// Send a request message to the RPC Server.  The response will be included in the IEnumerable
        /// returned by the GetIncomingSerializedMessages method.
        /// </summary>
        /// <param name="serializedMessage">A serialized RpcMessage.</param>
        /// <param name="requestTimeout">
        /// The timeout for this operation.Can only be set to System.Threading.Timeout.InfiniteTimeSpan or a TimeSpan value with total milliseconds greater or equal to 0 and less than 2147483648.
        /// </param>
        public void SendRequest(string serializedMessage, TimeSpan requestTimeout)
        {
            this.ThrowIfDisposed();

            if (!this.IsConnected)
            {
                throw new RpcException("Cannot send remote request until the transport has been connected.");
            }

            ThrowIfInvalidTimeout(requestTimeout, "requestTimeout");

            var socketRequest = new RpcSocketTransportMessage() { ServerId = this.serverId, Version = Assembly.GetAssembly(this.GetType()).GetName().Version.ToString(), RpcMessage = serializedMessage };

            string serializedSocketRequest;
            try
            {
                serializedSocketRequest = this.SerializeSocketTransportMessage(socketRequest);
            }
            catch (SerializationException ex)
            {
                throw new RpcException("Serialization of the RpcSocketTransportMessage failed. See inner exception for details.", ex);
            }

            try
            {
                this.WriteRequest(serializedSocketRequest, requestTimeout);
            }
            catch (IOException ex)
            {
                SocketException socketException = ex.InnerException as SocketException;
                if (socketException != null && socketException.SocketErrorCode == SocketError.TimedOut)
                {
                    throw new TimeoutException("Operation timed out. See inner exception for details.", ex);
                }

                throw new RpcException("Error while sending data. See inner exception for details.", ex);
            }

            this.SignalTestThatRpcMethodHasBeenCalled();
        }

        /// <summary>
        /// Connect to the target.
        /// </summary>
        /// <param name="timeout">The timeout to give to the connection attempt.</param>
        public void Connect(TimeSpan? timeout)
        {
            this.ThrowIfDisposed();

            if (this.IsConnected)
            {
                return;
            }

            // The connection is no longer valid, so clean up resources if they exist.
            this.CloseSocket();

            try
            {
                this.ConnectSocket(timeout);
            }
            catch (SocketException ex)
            {
                this.CloseSocket();

                throw new RpcException("Failed to connect.  See inner exception for details.", ex);
            }
        }

        /// <summary>
        /// Synchronously enumerates all incoming messages.
        /// This method will block until either the provided CancellationToken is canceled or the RpcSocketTransport is disposed.
        /// Therefore, it is recommended that users not call this method on the program's primary thread.
        /// Note: anything besides an RpcException thrown out of this function will cause unit tests to crash because of unhandled exception on background thread.
        /// </summary>
        /// <param name="cancellationToken">A cancellation token that can be used to terminate this operation.</param>
        /// <returns>A enumeration of all incoming serialized RpcMessages.</returns>
        public IEnumerable<string> GetIncomingSerializedMessages(CancellationToken cancellationToken)
        {
            this.ThrowIfDisposed();

            // Since there doesn't seem to be a good way to directly cancel a network stream,
            // we will register closing the underlying socket with the CancellationToken to
            // force the thread waiting on the stream to throw an exception
            cancellationToken.Register(() => this.CloseSocket());

            while (!cancellationToken.IsCancellationRequested)
            {
                string incomingString = null;

                try
                {
                    using (BinaryReader reader = new BinaryReader(this.NetworkStream, Encoding.UTF8, true))
                    {
                        RpcSocketTransportMessage transportMessage = this.DeserializeSocketTransportMessage(reader.ReadString());
                        RpcExceptionFactory.ThrowExceptionForRpcResponseCode(transportMessage.ErrorCode);

                        incomingString = transportMessage.RpcMessage;
                    }
                }
                catch (ObjectDisposedException)
                {
                    // This can happen when the network stream is disposed which
                    // occurs when this method is being cancelled.
                    yield break;
                }
                catch (IOException ex)
                {
                    // Handles any IOExceptions thrown by NetworkStream and Socket
                    if (this.IsDisconnectException(ex))
                    {
                        yield break;
                    }
                    else
                    {
                        throw new RpcException(IncomingSerializedMessagesFailure, ex);
                    }
                }
                catch (SerializationException ex)
                {
                    // Handles SerializationException thrown from DataContractJsonSerializer
                    throw new RpcException(IncomingSerializedMessagesFailure, ex);
                }
                catch (ArgumentException ex)
                {
                    // Handles ArgumentException and ArgumentNullException thrown from BinaryReader and MemoryStream
                    throw new RpcException(IncomingSerializedMessagesFailure, ex);
                }

                if (!string.IsNullOrWhiteSpace(incomingString))
                {
                    yield return incomingString;
                }
            }
        }

        /// <summary>
        /// Disconnect from the target.
        /// </summary>
        public void Disconnect()
        {
            this.ThrowIfDisposed();

            this.CloseSocket();
        }

        /// <summary>
        /// Used for test infrastructure.
        /// </summary>
        protected virtual void SignalTestThatRpcMethodHasBeenCalled()
        {
            // Please do not add any code here. If you need any functionality implement that in the override of the inherited class.
        }

        /// <summary>
        /// Disposes of the managed resources used by this class.
        /// </summary>
        protected override void DisposeManagedResources()
        {
            this.CloseSocket();
        }

        /// <summary>
        /// Creates and connects a socket to destination IP/port number.
        /// Also creates a network stream for reading from the socket.
        /// </summary>
        /// <param name="timeout">The timeout to give to the connection attempt.</param>
        private void ConnectSocket(TimeSpan? timeout)
        {
            try
            {
                this.socketLock.EnterWriteLock();

                this.socket = new Socket(SocketType.Stream, ProtocolType.Tcp);
                this.socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
                if (timeout.HasValue)
                {
                    var result = this.socket.BeginConnect(this.ipAddress, this.portNumber, null, null);

                    var succeeded = result.AsyncWaitHandle.WaitOne(timeout.Value);

                    if (succeeded || this.socket.Connected)
                    {
                        this.socket.EndConnect(result);
                    }
                    else
                    {
                        this.socket.Close();
                        this.socket = null;
                        throw new TimeoutException("Connection to remote host timed out.");
                    }
                }
                else
                {
                    this.socket.Connect(this.ipAddress, this.portNumber);
                }
            }
            finally
            {
                this.socketLock.ExitWriteLock();
            }

            lock (this.networkStreamLock)
            {
                // If the above code throws, the NetworkStream should not be created.
                // If the socket is not connected, the network stream constructor will throw.
                this.NetworkStream = new NetworkStream(this.socket, true);
            }
        }

        /// <summary>
        /// Disconnects and disposes the socket.
        /// </summary>
        private void CloseSocket()
        {
            lock (this.networkStreamLock)
            {
                try
                {
                    // Re-entry guard because this function is often called from random exception handlers, no single point.
                    this.socketLock.EnterWriteLock();

                    if (this.socket != null)
                    {
                        this.socket.Dispose();
                        this.socket = null;
                    }
                }
                finally
                {
                    this.socketLock.ExitWriteLock();
                }
            }
        }
        
        /// <summary>
        /// Determines whether the specified exception occured due to a socket disconnect.
        /// </summary>
        /// <param name="error">Exception object to examine.</param>
        /// <returns>Whether the specified exception occured due to socket disconnect.</returns>
        private bool IsDisconnectException(IOException error)
        {
            SocketException socketException = error.InnerException as SocketException;

            if (socketException != null)
            {
                if (/* This happens if the client chooses to disconnect or dispose of the channel. */
                    socketException.SocketErrorCode == SocketError.ConnectionAborted ||
                    /* This happens if the server shuts down while the client is connected. */
                    socketException.SocketErrorCode == SocketError.ConnectionReset)
                {
                    return true;
                }
            }

            ObjectDisposedException objectDisposedException = error.InnerException as ObjectDisposedException;

            if (objectDisposedException != null)
            {
                if (/* This happens when cancellation token is used to close the socket. */
                    objectDisposedException.ObjectName == typeof(System.Net.Sockets.Socket).ToString())
                {
                    return true;
                }
            }

            return false;
        }

        /// <summary>
        /// Serializes an RpcSocketTransportMessage to a string.
        /// </summary>
        /// <param name="request">The transport message object to serialize.</param>
        /// <returns>The serialized message as a string.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Usage", "CA2202:Do not dispose objects multiple times", Justification = "The StreamReader will not dispose of the stream.")]
        private string SerializeSocketTransportMessage(RpcSocketTransportMessage request)
        {
            using (MemoryStream stream = new MemoryStream())
            {
                this.socketTransportMessageSerializer.WriteObject(stream, request);
                stream.Position = 0;
                using (StreamReader reader = new StreamReader(stream))
                {
                    return reader.ReadToEnd();
                }
            }
        }

        /// <summary>
        /// Deserializes an RpcSocketTransportMessage from a string.
        /// </summary>
        /// <param name="serializedMessage">Serialized message to deserialize.</param>
        /// <returns>The transport message object.</returns>
        private RpcSocketTransportMessage DeserializeSocketTransportMessage(string serializedMessage)
        {
            try
            {
                using (MemoryStream stream = new MemoryStream(Encoding.UTF8.GetBytes(serializedMessage)))
                {
                    return (RpcSocketTransportMessage)this.socketTransportMessageSerializer.ReadObject(stream);
                }
            }
            catch (EncoderFallbackException ex)
            {
                // Handles EncoderFallbackException thrown by UTF8.GetBytes
                throw new RpcException("Failed to deserialize socket transport message", ex);
            }
        }

        /// <summary>
        /// Writes a serialized request to the network stream to get it sent to the server.
        /// </summary>
        /// <param name="serializedSocketRequest">The serialized Json request.</param>
        /// <param name="remainingTimeout">Remaining time left from the high-level timeout for sending the request.</param>
        private void WriteRequest(string serializedSocketRequest, TimeSpan remainingTimeout)
        {
            // Locking because only one thread can write to a network stream at a time.
            lock (this.networkStreamLock)
            {
                this.NetworkStream.WriteTimeout = checked((int)remainingTimeout.TotalMilliseconds);

                using (BinaryWriter writer = new BinaryWriter(this.NetworkStream, Encoding.UTF8, true))
                {
                    writer.Write(serializedSocketRequest);
                }
            }
        }        
    }
}
